import json

def build_and_save_traj(json_path, out_path=None):
    """
    读取场景 JSON，提取每个 TrajSegmentInfo 的 shape 和按时间划分的状态，
    构造出一个字典：
      {
        1: {
          "shape": { … },
          "0.0": { "x":…, "y":…, "yaw_rad":…, "v_mps":…, "yawrate_radps":…, "acc_mpss":… },
          "0.1": { … },
          …
        },
        2: { … }
      }
    如果 out_path != None，还会把这个结果写到 out_path 对应的文件（JSON）里。
    """
    # 1) 读 JSON
    with open(json_path, 'r', encoding='utf-8') as f:
        scene = json.load(f)

    dt = scene.get("dt", 0.1)

    traj = {}
    for idx, seg in enumerate(scene.get("TrajSegmentInfo", []), start=1):
        seg_dict = {}
        # 保留 shape 信息
        seg_dict["shape"] = seg["VehicleShapeInfo"]
        states = seg.get("states", {})

        # 假定 x,y,yaw_rad,v_mps,yawrate_radps,acc_mpss 都有相同长度
        length = len(states.get("x", []))
        for j in range(length):
            t = round(j * dt, 1)
            t_key = f"{t:.1f}"
            # 从每个 states[key] 中取出第 j 个元素的第 0 项
            seg_dict[t_key] = {
                "x":            float(states["x"][j][0]),
                "y":            float(states["y"][j][0]),
                "yaw_rad":      float(states["yaw_rad"][j][0]),
                "v_mps":        float(states["v_mps"][j][0]),
                "yawrate_radps":float(states["yawrate_radps"][j][0]),
                "acc_mpss":     float(states["acc_mpss"][j][0]),
            }

        traj[idx] = seg_dict

    # 3) 如果指定了输出路径，把 traj 写成 JSON
    if out_path:
        with open(out_path, 'w', encoding='utf-8') as f:
            json.dump(traj, f, ensure_ascii=False, indent=2)

    return traj

# ===== 用法示例 =====
if __name__ == "__main__":
    # 1) 读原始场景 JSON，生成 traj 变量
    traj = build_and_save_traj("scene.json",
                               out_path="traj_extracted.json")

    # 2) 现在 traj 看起来就像你要的格式
    #    比如 traj[1]["shape"]、traj[1]["0.0"]、traj[1]["0.1"] ...
    print("vehicles:", list(traj.keys()))
    print("first vehicle shape:", traj[1]["shape"])
    print("t=0.0 state:",       traj[1]["0.0"])
    print("t=0.1 state:",       traj[1]["0.1"])